#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : wlisingle
# @Email : 19688513@qq.com
# @Project : g-carbon-bio
# @File : stock_data_fetcher.py
# @Software: PyCharm
# @time: 2025/3/17 10:29
# @description : 一个用于获取沪 A 股的实时行情数据，另一个用于获取每只股票的资金流数据并合并输出。
"""

import akshare as ak
import pandas as pd
from datetime import datetime, timedelta


class StockDataFetcher:
    def __init__(self):
        pass

    def fetch_sh_a_spot_data(self):
        """
        获取沪 A 股实时行情数据
        :return: DataFrame, 包含沪 A 股实时行情数据
        """
        return ak.stock_sh_a_spot_em()

    def fetch_sz_a_spot_data(self):
        """
        获取深 A 股实时行情数据
        :return: DataFrame, 包含深 A 股实时行情数据
        """
        return ak.stock_sz_a_spot_em()

    def fetch_fund_flow_rank(self, indicator="今日"):
        """
        获取个股资金流排名数据
        :param indicator: str, 排名时间范围（"今日", "3日", "5日", "10日"）
        :return: DataFrame, 包含个股资金流排名数据
        """
        return ak.stock_individual_fund_flow_rank(indicator=indicator)

    def fetch_historical_data(self, symbol, start_date, end_date, adjust="qfq"):
        """
        获取指定股票的历史行情数据
        :param symbol: str, 股票代码（如 "sh600000"）
        :param start_date: str, 开始日期（如 "20230101"）
        :param end_date: str, 结束日期（如 "20231031"）
        :param adjust: str, 复权类型（"qfq": 前复权, "hfq": 后复权）
        :return: DataFrame, 包含历史行情数据
        """
        return ak.stock_zh_a_daily(symbol=symbol, start_date=start_date, end_date=end_date, adjust=adjust)

    def calculate_moving_averages(self, historical_data):
        """
        计算均线（5日、10日、20日、30日、60日、120日）
        :param historical_data: DataFrame, 包含历史行情数据
        :return: DataFrame, 包含均线数据
        """
        historical_data['5日均线'] = historical_data['close'].rolling(window=5).mean()
        historical_data['10日均线'] = historical_data['close'].rolling(window=10).mean()
        historical_data['20日均线'] = historical_data['close'].rolling(window=20).mean()
        historical_data['30日均线'] = historical_data['close'].rolling(window=30).mean()
        historical_data['60日均线'] = historical_data['close'].rolling(window=60).mean()
        historical_data['120日均线'] = historical_data['close'].rolling(window=120).mean()
        return historical_data

    def fetch_combined_data(self, market='sh', indicator="今日", enable_history=True):
        """
        获取指定市场（沪市或深市）的实时行情数据，并合并个股资金流排名数据
        :param market: str, 市场代码（'sh' 表示沪市，'sz' 表示深市）
        :param indicator: str, 排名时间范围（"今日", "3日", "5日", "10日"）
        :param enable_history: bool, 是否读取历史数据计算均线（默认True）
        :return: DataFrame, 包含实时行情和资金流排名数据的合并结果
        """
        # 根据市场代码获取实时行情数据
        if market == 'sh':
            spot_data = self.fetch_sh_a_spot_data()
        elif market == 'sz':
            spot_data = self.fetch_sz_a_spot_data()
        else:
            raise ValueError("市场代码错误！仅支持 'sh'（沪市）或 'sz'（深市）。")

        # 获取个股资金流排名数据
        fund_flow_rank_data = self.fetch_fund_flow_rank(indicator=indicator)

        # 创建一个空的 DataFrame 用于存储最终结果
        result_df = pd.DataFrame()

        # 获取当前日期并往前推 150 天（仅当启用历史数据时计算）
        if enable_history:
            end_date = datetime.now().strftime("%Y%m%d")
            start_date = (datetime.now() - timedelta(days=150)).strftime("%Y%m%d")

        # 遍历每一只股票
        for index, row in spot_data.iterrows():
            stock_code = row['代码']
            print(f"正在处理股票: {stock_code}")

            try:
                # 从资金流排名数据中筛选出当前股票的数据
                stock_fund_flow = fund_flow_rank_data[fund_flow_rank_data['代码'] == stock_code]

                # 如果找到匹配的资金流数据
                if not stock_fund_flow.empty:
                    # 合并基础数据
                    combined_row = pd.concat([row, stock_fund_flow.iloc[0]])

                    # 处理历史数据和均线（仅在启用时执行）
                    if enable_history:
                        symbol = f"{market}{stock_code}"
                        historical_data = self.fetch_historical_data(symbol, start_date, end_date, adjust="qfq")
                        if not historical_data.empty:
                            historical_data_with_ma = self.calculate_moving_averages(historical_data)
                            latest_ma = historical_data_with_ma.iloc[-1][
                                ['5日均线', '10日均线', '20日均线', '30日均线', '60日均线', '120日均线']]
                            combined_row = pd.concat([combined_row, latest_ma])
                        else:
                            print(f"警告: 股票 {stock_code} 无历史数据")

                    # 添加到结果集
                    result_df = pd.concat([result_df, combined_row.to_frame().T], ignore_index=True)
            except Exception as e:
                print(f"股票 {stock_code} 数据处理失败: {e}")

        return result_df


# 使用示例
if __name__ == "__main__":
    fetcher = StockDataFetcher()

    # 控制开关设置（True启用历史数据计算均线，False跳过）
    ENABLE_HISTORY = False  # 在此处修改开关状态

    # 获取沪 A 股数据（不计算均线）
    print("获取沪 A 股数据...")
    combined_data_sh = fetcher.fetch_combined_data(market='sh', indicator="今日", enable_history=ENABLE_HISTORY)
    print(combined_data_sh)

    # 获取深 A 股数据（不计算均线）
    print("获取深 A 股数据...")
    combined_data_sz = fetcher.fetch_combined_data(market='sz', indicator="今日", enable_history=ENABLE_HISTORY)
    print(combined_data_sz)

    # 生成带时间戳的文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"stock_data_{timestamp}.xlsx"

    # 将数据保存到 Excel 文件
    with pd.ExcelWriter(filename) as writer:
        combined_data_sh.to_excel(writer, sheet_name='沪 A 股', index=False)
        combined_data_sz.to_excel(writer, sheet_name='深 A 股', index=False)

    print(f"数据已保存到 {filename}")